import aiohttp
import asyncio
import secrets
import random
from typing import Dict, Any, List, Tuple
import json
from .flux_config import get_flux_payload

class FluxAPIService:
    def __init__(self):
        self.base_url = "https://u142199-b6d6-39157cb1.westx.seetacloud.com:8443"
        self.headers = {
            'Content-Type': 'application/json',
        }
    
    async def submit_prompt(self, image_url: str, prompt: str = "", variation_count: int = 4) -> Tuple[bool, Dict[str, Any]]:
        """提交Flux生成任务"""
        try:
            # 生成随机client_id和seed
            client_id = secrets.token_hex(16)
            seed = random.randint(10 ** 13, 10 ** 14 - 1)
            
            payload = get_flux_payload(client_id, prompt, image_url, variation_count, seed)
            
            print(f"✅ 提交Flux任务，client_id: {client_id}")
            
            # 配置连接器以解决SSL和连接管理问题
            connector = aiohttp.TCPConnector(
                limit=20,
                limit_per_host=10,
                keepalive_timeout=60,
                enable_cleanup_closed=True,
                ssl=False  # 禁用SSL验证以避免SSL握手问题
            )
            async with aiohttp.ClientSession(connector=connector, timeout=aiohttp.ClientTimeout(total=120)) as session:
                async with session.post(
                    f"{self.base_url}/prompt",
                    headers=self.headers,
                    json=payload,
                    timeout=aiohttp.ClientTimeout(total=60)
                ) as response:
                    if response.status == 200:
                        result = await response.json()
                        
                        # 尝试多种可能的任务ID字段
                        task_id = None
                        possible_keys = ['id', 'task_id', 'prompt_id', 'client_id']
                        for key in possible_keys:
                            if key in result:
                                task_id = result[key]
                                break
                        
                        if task_id:
                            print(f"✅ Flux任务提交成功，任务ID: {task_id}")
                            return True, {'task_id': task_id, 'full_response': result}
                        else:
                            error_msg = f"响应中未找到任务ID，可用字段: {list(result.keys())}"
                            print(f"❌ {error_msg}")
                            return False, {"error": error_msg}
                    else:
                        response_text = await response.text()
                        error_msg = f"Flux API请求失败: HTTP {response.status}: {response_text}"
                        print(f"❌ {error_msg}")
                        return False, {"error": error_msg}
                        
        except asyncio.TimeoutError:
            error_msg = "Flux API请求超时"
            return False, {"error": error_msg}
        except Exception as e:
            error_msg = f"Flux API请求异常: {str(e)}"
            return False, {"error": error_msg}
    
    async def get_task_history(self, task_id: str) -> Tuple[bool, Dict[str, Any]]:
        """获取任务历史和状态"""
        try:
            # 配置连接器以解决SSL和连接管理问题
            connector = aiohttp.TCPConnector(
                limit=20,
                limit_per_host=10,
                keepalive_timeout=60,
                enable_cleanup_closed=True,
                ssl=False  # 禁用SSL验证以避免SSL握手问题
            )
            async with aiohttp.ClientSession(connector=connector, timeout=aiohttp.ClientTimeout(total=60)) as session:
                async with session.get(
                    f"{self.base_url}/history/{task_id}",
                    headers=self.headers,
                    timeout=aiohttp.ClientTimeout(total=30)
                ) as response:
                    if response.status == 200:
                        result = await response.json()
                        return True, result
                    else:
                        response_text = await response.text()
                        return False, {"error": f"HTTP {response.status}: {response_text}"}
                        
        except Exception as e:
            return False, {"error": f"Request failed: {str(e)}"}
    
    async def get_image_view(self, output: str, filename: str) -> Tuple[bool, bytes]:
        """获取最终生成的图片"""
        try:
            params = {
                'output': output,
                'filename': filename
            }
            
            # 配置连接器以解决SSL和连接管理问题
            connector = aiohttp.TCPConnector(
                limit=20,
                limit_per_host=10,
                keepalive_timeout=60,
                enable_cleanup_closed=True,
                ssl=False  # 禁用SSL验证以避免SSL握手问题
            )
            async with aiohttp.ClientSession(connector=connector, timeout=aiohttp.ClientTimeout(total=120)) as session:
                async with session.get(
                    f"{self.base_url}/view",
                    params=params,
                    timeout=aiohttp.ClientTimeout(total=60)
                ) as response:
                    if response.status == 200:
                        image_data = await response.read()
                        print(f"✅ 成功获取图片: {filename}")
                        return True, image_data
                    else:
                        response_text = await response.text()
                        return False, f"HTTP {response.status}: {response_text}".encode()
                        
        except Exception as e:
            return False, f"Request failed: {str(e)}".encode()
    
    async def wait_for_completion_and_get_images(self, task_id: str, max_wait_time: int = 300) -> Tuple[bool, List[bytes]]:
        """等待任务完成并获取所有生成的图片"""
        start_time = asyncio.get_event_loop().time()
        
        while asyncio.get_event_loop().time() - start_time < max_wait_time:
            # 获取任务历史
            success, result = await self.get_task_history(task_id)
            
            if not success:
                await asyncio.sleep(10)
                continue
            
            # 检查响应格式：{task_id: {outputs: {...}, status: {...}}}
            task_data = None
            if isinstance(result, dict):
                # 查找任务数据，可能在task_id键下或直接在根级别
                if task_id in result:
                    task_data = result[task_id]
                elif 'outputs' in result:
                    task_data = result
            
            if not task_data:
                await asyncio.sleep(10)
                continue
            
            # 检查任务状态
            status_info = task_data.get('status', {})
            if isinstance(status_info, dict):
                status_str = status_info.get('status_str')
                completed = status_info.get('completed', False)
                
                if status_str == 'success' and completed:
                    print(f"✅ 任务状态: {status_str}")
                elif status_str == 'error':
                    print(f"❌ 任务失败: {status_str}")
                    return False, []
            
            # 查找outputs字段中的图片信息
            images = []
            outputs = task_data.get('outputs', {})
            
            if isinstance(outputs, dict):
                # 遍历所有输出节点
                for node_id, node_data in outputs.items():
                    if isinstance(node_data, dict) and 'images' in node_data:
                        image_list = node_data['images']
                        
                        if isinstance(image_list, list):
                            for img_info in image_list:
                                if isinstance(img_info, dict):
                                    filename = img_info.get('filename')
                                    img_type = img_info.get('type')
                                    
                                    if filename and img_type == 'output':
                                        print(f"✅ 发现图片: {filename}")
                                        
                                        # 使用节点ID作为output参数
                                        img_success, img_data = await self.get_image_view(node_id, filename)
                                        
                                        if img_success and isinstance(img_data, bytes):
                                            images.append(img_data)
                                        else:
                                            print(f"❌ 获取图片失败: {filename}")
            
            if images:
                print(f"✅ 任务完成，获取到 {len(images)} 张图片")
                return True, images
            
            await asyncio.sleep(10)  # 等待10秒后重试
        
        print(f"❌ 任务超时")
        return False, []
    
    async def generate_variations(self, image_url: str, prompt: str = "", variation_count: int = 4) -> Tuple[bool, List[bytes]]:
        """完整的Flux变体生成流程"""
        try:
            # 第一步：提交任务
            success, submit_result = await self.submit_prompt(image_url, prompt, variation_count)
            
            if not success:
                error_msg = f"Flux任务提交失败: {submit_result.get('error', '未知错误')}"
                print(f"❌ {error_msg}")
                return False, []
            
            task_id = submit_result['task_id']
            print(f"开始等待任务完成...")
            
            # 第二步：等待完成并获取图片
            success, images = await self.wait_for_completion_and_get_images(task_id)
            
            if success:
                print(f"✅ Flux任务完成，生成图片数量: {len(images)}")
                return True, images
            else:
                print(f"❌ Flux任务失败")
                return False, []
                
        except Exception as e:
            error_msg = f"Flux生成过程异常: {str(e)}"
            print(f"❌ {error_msg}")
            return False, []

flux_service = FluxAPIService()